#!/usr/bin/env python3
# coding:utf-8
import os
import sys

sys.path.insert(0, "/home/ucar/ucar_ws/src/darknet_pkg/src/darknet_yolov4_tiny")
from darknet_yolov4_tiny.Detect import Detect
import cv2
from GetCvPicture import PicGetter
from geometry_msgs.msg import Twist
import rospy
from vision.srv import darknet_message, darknet_messageRequest, darknet_messageResponse
import time
import numpy as np
import threading

class Detector:
    def __init__(self):
        # Darknet目标检测实体对象
        self.detect = Detect()
        self.base_conf = 0.8
        self.task = ""

    def Predict(self, img):
        res, plot = self.detect.predict_image(img, self.base_conf)

        # 使用zip进行解包，然后使用map转换为list
        # name, conf, xywh = map(list, zip(*res))

        if self.task == "水果":
            for result in res:
                 if result[0] in ["watermelon", "banana", "apple"]:
                    return result

        elif self.task == "蔬菜":
            for result in res:
                if result[0] in ["tomato", "potato", "pepper"]:
                    return result
        else:
            for result in res:
                if result[0] in ["coke", "milk", "cake"]:
                    return result
        return None

    def SetTask(self, task):
        self.task = task


class Controller:
    def __init__(self):
        self.KP = 2.8
        self.max_output = 0.2
        '''
        RGB 相机参数 
        # 单目超广角 彩色图像分辨率：19201080@30FPS、1280720@30FPS、640*480@30FPS 
        # 视场角：水平FOV124.8°、垂直FOV67°、对角FOV160° 
        # 焦距：2.8mm F/NO(Infinite)：2.6±5% 
        # 物距：45cm—100m
        '''
        self.image_width = 640
        self.image_height = 480
        self.FOV_width = 124.8  # 相机水平方向视场角

    def Revise(self, xywh):
        ''' 更新偏差角'''
        # yaw轴姿态偏移(弧度) = （（目标相机位置 - 相机图片宽度）/ 相机图片宽度 ） * 相机水平方向视场角 / 180 * 3.14
        yaw_diff = (((xywh[0] - self.image_width / 2) / self.image_width) * self.FOV_width / 180) * 3.1419

        if abs(yaw_diff) <= 0.04:
            return None

        else:
            output = self.KP * yaw_diff
            output = np.clip(output, -self.max_output, self.max_output)
            return output



class Detect_Service:
    def __init__(self):
        self.cap = PicGetter()
        self.cmd_pub = rospy.Publisher("/cmd_vel", Twist, queue_size=10)
        self.detector = Detector()
        self.service = None
        self.controller = Controller()

        # Speed publish loop
        self.current_twist = Twist()
        self.publish_rate = 20  # Hz
        self.running = False
        self.publisher_thread = threading.Thread(target=self.publish_loop)
        self.publisher_thread.daemon = True

    def publish_loop(self):
        rospy.loginfo("Publish speed")
        """单独线程持续发布速度"""
        rate = rospy.Rate(self.publish_rate)
        print("self.running:",self.running)
        while not rospy.is_shutdown():
            if self.running:
                self.cmd_pub.publish(self.current_twist)
                print("send speed")
            rate.sleep()

    def doRequest(self, req):
        # 进行一次操作
        task = req.task
        rospy.loginfo(f"本次采购任务为：{task}")
        self.detector.SetTask(task)

        self.running = True
        # 设置旋转速度
        self.current_twist = Twist()
        self.current_twist.angular.z = 0.7

        res = None

        rate = rospy.Rate(25)
        start = time.time()

        while time.time() - start < 10:
            img = self.cap.Get()
            if img is None:
                continue
            res = self.detector.Predict(img)

            if res is not None:
                self.current_twist.angular.z = 0
                break
        if res is None:
            self.running = False
            return darknet_messageResponse(False, "未检测到需要采集的物品")
        
        # 接下来就是要转动或者平移，直到中线点与识别点对齐
        rospy.loginfo(f"我已发现物品{res[0]},马上开始追踪")

        count = 0

        while 1:
            if res is None:
                rospy.loginfo("目标暂时丢失")
                img = self.cap.Get()
                if img is None:
                    continue
                res = self.detector.Predict(img)
                continue
            output = self.controller.Revise(res[2])
            rospy.loginfo(f"output: {output}")
            count += 1
            if output is None:
                rospy.loginfo("已经朝向既定目标")
                self.current_twist.angular.z = 0
                self.running = False
                return darknet_messageResponse(True, res[0])
            
            self.current_twist.angular.z = -output
            
            if count >= 50:
                rospy.loginfo("已经多次矫正")
                self.current_twist.angular.z = 0
                self.running = False
                return darknet_messageResponse(True, res[0])
            res = self.detector.Predict(self.cap.Get())

        

    def start(self):
        rospy.init_node('DarknetNode', anonymous=True)
        self.cap.start()
        self.publisher_thread.start()
        rospy.loginfo("Darknet_Detect Service is ready")
        self.service = rospy.Service("Darknet_DetectService", darknet_message, self.doRequest)


if __name__ == '__main__':
    service = Detect_Service()

    service.start()

    rospy.spin()
